package org.finesys.common.log.trace.tracer;

import com.alibaba.ttl.TransmittableThreadLocal;
import lombok.Data;
import org.finesys.common.log.trace.util.TraceIdUtil;
import org.slf4j.MDC;
import org.springframework.util.StringUtils;

import static org.finesys.common.log.constants.TraceConstants.*;

@Data
public final class Tracer {


    private static TransmittableThreadLocal<Tracer> traceThreadLocal = new TransmittableThreadLocal<>();
    private String traceId;
    private String spanId;
    private String parentId;

    public Tracer(TracerBuilder builder) {
        this.traceId = builder.traceId;
        this.spanId = builder.spanId;
        this.parentId = builder.parentId;
        traceThreadLocal.set(this);
    }


    /**
     * trace root
     *
     * @param traceId  the trace id
     * @param spanId   the span id
     * @param parentId the parent id
     */
    public static void trace(String traceId, String spanId, String parentId) {
        Tracer tracer = Tracer.traceThreadLocal.get();
        if (null == tracer) {
            tracer = new TracerBuilder().traceId(traceId).spanId(spanId).parentId(parentId).build();
            tracer.buildSpan();
        } else {
            traceThreadLocal.remove();
        }
    }

    /**
     * build span
     */
    public void buildSpan() {
        if (!StringUtils.hasText(traceId)) {
            traceId = TraceIdUtil.traceId();
        }
        if (!StringUtils.hasText(parentId)) {
            parentId = "-1";
        }
        if (!StringUtils.hasText(spanId)) {
            spanId = TraceIdUtil.spanId();
        }
        parentId = TraceIdUtil.add(parentId);
        spanId = TraceIdUtil.add(spanId);
        MDC.put(TRACE_ID, traceId);
        MDC.put(SPAN_ID, spanId);
        MDC.put(PARENT_ID, parentId);
    }

    public static void remove() {
        traceThreadLocal.remove();
        MDC.clear();
    }

    public static class TracerBuilder {
        private String traceId;
        private String spanId;
        private String parentId;

        /**
         * build trace id
         *
         * @param traceId trace id
         * @return trace id
         */
        public TracerBuilder traceId(String traceId) {
            this.traceId = traceId;
            return this;
        }

        /**
         * build span id
         *
         * @param spanId span id
         * @return span id
         */
        public TracerBuilder spanId(String spanId) {
            this.spanId = spanId;
            return this;
        }

        /**
         * build parent id
         *
         * @param parentId parent id
         * @return parent id
         */
        public TracerBuilder parentId(String parentId) {
            this.parentId = parentId;
            return this;
        }

        /**
         * build
         *
         * @return Tracer
         */
        public Tracer build() {
            return new Tracer(this);
        }
    }
}
